import pandas as pd
from . import config, utils
import logging
import pickle


def main():
    logging.warning('Will set task, dataset = test')
    config.task = 'test'
    config.dataset = 'test'

    test = pd.read_csv(config.pj_root + 'data/' + config.dataset + '.csv', index_col='no')
    a_feature = pd.read_csv(config.pj_root + 'data/a_feature_test.csv', index_col='no')
    test = test.join(a_feature)

    if len(config.drop_columns) != 0:
        logging.warning('Will drop %s len=%d' % (str(config.drop_columns), len(config.drop_columns)))
        test = test.drop(config.drop_columns, axis=1)

    if len(config.select_columns) != 0:
        logging.warning('Will select %s len=%d' % (str(config.select_columns), len(config.select_columns)))
        test = test[test.columns[test.columns.isin(config.select_columns)]]
    else:
        config.select_columns = list(test.columns)[0:]

    if config.use_basic_process:
        test = utils.basic_process_use_map(test, has_flag=False)

    X = test.values[:, 0:]  # all columns

    with open(config.pj_root + 'model/%s.mo' % config.model.__name__, 'rb') as f:
        model = pickle.load(f)
    pred = model.predict_proba(X)

    pred = pred[:, 1]
    pred_df = pd.DataFrame(index=test.index)
    pred_df = pred_df.join(pd.DataFrame(pred, index=test.index, columns=['pred']))
    pred_df.to_csv(config.pj_root + 'result/qianhai_%s.csv' % (config.model.__name__))

if __name__ == '__main__':
    main()
